# -*- coding: utf-8 -*-  
'''
nlu模型零件

Created on 2021年10月2日
@author: luoyi
'''
import tensorflow as tf

from utils.relationship import Relationship


#    意图layer
class InformLayer(tf.keras.layers.Layer):
    def __init__(self,
                 name='inform_layer',
                 inform_size=Relationship.instance().rel_size(),
                 **kwargs):
        self._name = name
        self._inform_size = inform_size
        
        super(InformLayer, self).__init__(name=name, **kwargs)
        pass
    
    def build(self, input_shape):
        #    意图分类
        self._fc = tf.keras.layers.Dense(name=self._name + '_fc', 
                                         units=self._inform_size, 
                                         kernel_initializer=tf.initializers.he_normal(), 
                                         bias_initializer=tf.initializers.he_normal())
        super(InformLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        '''
            @param inputs: [CLS]位的语义向量    Tensor(batch_size, inform_size)
        '''
        out = self._fc(inputs)
        out = tf.nn.softmax(out)
        self._out = out
        return out
    
    #    拿fc层输出
    def get_out(self):
        return self._out
    pass
